import * as tf from "@tensorflow/tfjs-node";

// 生成随机训练数据
function generateTrainingData(count) {
  const data = [];
  for (let i = 0; i < count; i++) {
    const input = Array(33)
      .fill(0)
      .map(() => Math.floor(Math.random() * 15));
    const output = Array(6)
      .fill(0)
      .map(() => Math.random());
    data.push({ input, output });
  }
  return data;
}

// 将数据转换为张量
function convertToTensors(data) {
  return tf.tidy(() => {
    tf.util.shuffle(data);

    const inputs = data.map((d) => d.input);
    const outputs = data.map((d) => d.output);

    const inputTensor = tf.tensor2d(inputs, [inputs.length, 33]);
    const outputTensor = tf.tensor2d(outputs, [outputs.length, 6]);

    const inputMax = inputTensor.max();
    const inputMin = inputTensor.min();
    const outputMax = outputTensor.max();
    const outputMin = outputTensor.min();

    // 归一化数据
    const normalizedInputs = inputTensor
      .sub(inputMin)
      .div(inputMax.sub(inputMin));
    const normalizedOutputs = outputTensor
      .sub(outputMin)
      .div(outputMax.sub(outputMin));

    return {
      inputs: normalizedInputs,
      outputs: normalizedOutputs,
      inputMax,
      inputMin,
      outputMax,
      outputMin,
    };
  });
}

// 创建模型
function createModel() {
  const model = tf.sequential();

  model.add(
    tf.layers.dense({ inputShape: [33], units: 64, activation: "relu" })
  );
  model.add(tf.layers.dense({ units: 32, activation: "relu" }));
  model.add(tf.layers.dense({ units: 16, activation: "relu" }));
  model.add(tf.layers.dense({ units: 6, activation: "sigmoid" }));

  model.compile({
    optimizer: tf.train.adam(0.001),
    loss: "meanSquaredError",
  });

  return model;
}

// 训练模型
async function trainModel(model, inputs, outputs) {
  return await model.fit(inputs, outputs, {
    batchSize: 32,
    epochs: 100,
    shuffle: true,
    validationSplit: 0.1,
    callbacks: tf.callbacks.earlyStopping({ monitor: "val_loss", patience: 5 }),
  });
}

// 预测函数
function predict(model, input, inputMax, inputMin, outputMax, outputMin) {
  const normalizedInput = tf
    .tensor2d([input])
    .sub(inputMin)
    .div(inputMax.sub(inputMin));
  const normalizedPrediction = model.predict(normalizedInput);
  const denormalizedPrediction = normalizedPrediction
    .mul(outputMax.sub(outputMin))
    .add(outputMin);
  return denormalizedPrediction.dataSync().map((n) => Math.round(n * 32) + 1);
}

// 主函数
async function run() {
  // 生成训练数据
  const trainingData = generateTrainingData(1000);

  // 转换数据为张量
  const { inputs, outputs, inputMax, inputMin, outputMax, outputMin } =
    convertToTensors(trainingData);

  // 创建模型
  const model = createModel();

  // 训练模型
  console.log("开始训练模型...");
  const history = await trainModel(model, inputs, outputs);
  console.log("训练完成");

  // 示例输入数据
  const exampleInput = [
    4, 4, 3, 1, 11, 11, 2, 3, 14, 4, 8, 2, 14, 5, 18, 7, 1, 3, 1, 1, 5, 4, 8, 2,
    1, 12, 7, 6, 43, 3, 7, 1, 2,
  ];

  // 使用训练好的模型进行预测
  const prediction = predict(
    model,
    exampleInput,
    inputMax,
    inputMin,
    outputMax,
    outputMin
  );
  console.log("预测结果:", prediction);

  // 测试模型
  const testData = generateTrainingData(10);
  console.log("\n测试模型:");
  testData.forEach((data, index) => {
    const testPrediction = predict(
      model,
      data.input,
      inputMax,
      inputMin,
      outputMax,
      outputMin
    );
    console.log(`测试 ${index + 1}:`);
    console.log(
      "实际输出:",
      data.output.map((n) => Math.round(n * 32) + 1)
    );
    console.log("预测输出:", testPrediction);
    console.log("---");
  });
}

run();
